"""
Preprocess dataset for countdown task - given a target number and N numbers, generate equations to reach target
"""

import re
import os
from datasets import Dataset, load_dataset
import random
from typing import List, Tuple
from tqdm import tqdm
from verl.utils.hdfs_io import copy, makedirs
import argparse
import pandas as pd
import json
import dotenv
dotenv.load_dotenv("/llm/nankai/xuyang_space/project/.env")
from openai import OpenAI


prompt_template = """A conversation involving User, Agents, and Decision - Maker. The User asks a question, and the Agents and the Decision - Maker work together to solve it. Each Agent provides an response to this question. These answers may contain correct parts, as well as incorrect, ambiguous, or irrelevant information. First, the Decision - Maker analyzes each agent's answer in the mind. After comprehensive analysis, the Decision - Maker makes a final judgment (which may adopt the strengths of multiple agents or completely reject their viewpoints), and finally provides the User with the answer. The analysis process and judgment are enclosed within <think> </think> and answer is enclosed within <answer> </answer> tags, respectively, i.e., <think> analysis process and judgment here </think> <answer> answer here </answer>.
User: {question}
Agent 1: {answer_1}
Agent 2: {answer_2}
Agent 3: {answer_3}
Decision - Maker: <think>
"""

data_source = 'agent_cdm'

local_dir = "/llm/nankai/xuyang_space/data/AM-DeepSeek-Distilled-40M/agent_cdm/raw/data20250514_extract_sft_rl_all_right_0610/rl/raw"
save_dir = "/llm/nankai/xuyang_space/data/AM-DeepSeek-Distilled-40M/agent_cdm/processed/data202506010"



train_dataset = load_dataset('json', data_files=f'{local_dir}/train.jsonl')
test_dataset = load_dataset('json', data_files=f'{local_dir}/test.jsonl')


train_dataset = train_dataset["train"]
test_dataset = test_dataset["train"]


def make_map_fn(split):
    def process_fn(example, idx):
        question = example["question"]
        solution = example["answer"]
        idx_list = [1,2,3]
        random.shuffle(idx_list)

        agent1 = example[f"agent{idx_list[0]}"]
        agent2 = example[f"agent{idx_list[1]}"]
        agent3 = example[f"agent{idx_list[2]}"]

        question_content = prompt_template.format(
            question=question,
            answer_1=agent1,
            answer_2=agent2,
            answer_3=agent3
        )
        data = {
            "data_source": data_source,
            "prompt": [{
                "role": "user",
                "content": question_content,
            }],
            "ability": "math",
            "reward_model": {
                "style": "rule",
                "ground_truth": solution
            },
            "extra_info": {
                'split': split,
                'index': idx,
            }
        }
        return data
    return process_fn

train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)
test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)


train_dataset.to_parquet(os.path.join(save_dir, 'train2.parquet'))
test_dataset.to_parquet(os.path.join(save_dir, 'test2.parquet'))

